{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "uRLPr0TnIAHO",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "BRANCH = 'r1.0.0rc1'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "o_0K1lsW1dj9"
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "You can run either this notebook locally (if you have all the dependencies and a GPU) or on Google Colab.\n",
    "\n",
    "Instructions for setting up Colab are as follows:\n",
    "1. Open a new Python 3 notebook.\n",
    "2. Import this notebook from GitHub (File -> Upload Notebook -> \"GITHUB\" tab -> copy/paste GitHub URL)\n",
    "3. Connect to an instance with a GPU (Runtime -> Change runtime type -> select \"GPU\" for hardware accelerator)\n",
    "4. Run this cell to set up dependencies.\n",
    "\"\"\"\n",
    "# If you're using Google Colab and not running locally, run this cell\n",
    "\n",
    "# install NeMo\n",
    "!python -m pip install git+https://github.com/NVIDIA/NeMo.git@$BRANCH#egg=nemo_toolkit[nlp]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# If you're not using Colab, you might need to upgrade jupyter notebook to avoid the following error:\n",
    "# 'ImportError: IProgress not found. Please update jupyter and ipywidgets.'\n",
    "\n",
    "! pip install ipywidgets\n",
    "! jupyter nbextension enable --py widgetsnbextension\n",
    "\n",
    "# Please restart the kernel after running this cell"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "dzqD2WDFOIN-"
   },
   "outputs": [],
   "source": [
    "from nemo.collections import nlp as nemo_nlp\n",
    "from nemo.utils.exp_manager import exp_manager\n",
    "\n",
    "import os\n",
    "import wget \n",
    "import torch\n",
    "import pytorch_lightning as pl\n",
    "from omegaconf import OmegaConf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "daYw_Xll2ZR9"
   },
   "source": [
    "In this tutorial, we are going to describe how to finetune a BERT-like model based on [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://arxiv.org/abs/1810.04805) on information retrieval task. We will be working on MSMARCO dataset which contains more than 800K queries and more than 8.8M short passages and the task is to rank all passages by their relevance to the particular query. Specifically, we will be training two different models.\n",
    "\n",
    "### BERT joint re-ranking \n",
    "\n",
    "* Original paper: [Passage Re-ranking with BERT](https://arxiv.org/abs/1901.04085).\n",
    "* Model overview: input query-passage pair is encoded as **[CLS] query_tokens [SEP] passage_tokens [SEP]** and is fed into BERT encoder. Last hidden state of [CLS] token is then passed into fully-connected layer to get similarity score.\n",
    "* Pros: high accuracy.\n",
    "* Cons: the model is too computationally demanding to use for all passages re-ranking. It is better to use this model for re-ranking the shortlist of ~top-100 candidates.\n",
    "\n",
    "### Dense Passage Retrieval\n",
    "\n",
    "* Original paper: [Dense Passage Retrieval for Open-Domain Question Answering](https://arxiv.org/abs/2004.04906).\n",
    "* Model overview: input query-passage pair is separately as **[CLS] query_tokens [SEP]** and **[CLS] passage_tokens [SEP]** which are fed into two different BERT encoders. Last hidden states of corresponding [CLS] tokens are treated as query and passage embedding respectively. The similarity score is computed as a dot-product between query and passage embeddings.\n",
    "* Pros: as the computation of query and passage embeddings is disentangled, all passage embeddings can be pre-computed in a single run through the passage collection. Then, we can build FAISS index on top of them and retrieve relevant passages very fast.\n",
    "* Cons: accuracy is lower comparing to joint model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ZnuziSwJ1yEB"
   },
   "source": [
    "# Dataset\n",
    "\n",
    "First of all, we need to download and prepare training dataset. Navigate to [examples/nlp/information_retrieval](https://github.com/NVIDIA/NeMo/blob/main/examples/nlp/information_retrieval) and download MSMARCO dataset:\n",
    "\n",
    "`bash get_msmarco.sh`\n",
    "\n",
    "The dataset should contain the following files:\n",
    "* collection.tsv - training passages with entries (passage_id, passage_text)\n",
    "* queries.(train/dev).tsv - training/development queries with entries (query_id, query_text)\n",
    "* qrels.(train/dev).tsv - training/development relevance scores with entries (query_id, 0, passage_id, 1)\n",
    "\n",
    "Information retrieval models are usually trained on lists of query, corresponding relevant passage, and several irrelevant passages in a contrastive manner. Thus, the format of training dataset for NeMo information retrieval models is the following:\n",
    "* collection.tsv.(pkl/npz) - tokenized and cached training passages\n",
    "* queries.train.tsv.(pkl/npz) - tokenized and cached training queries\n",
    "* query2passages.tsv - file with entries (query_id, rel_psg_id, irrel_psg_1_id, ..., irrel_psg_k_id)\n",
    "\n",
    "Note, the way we choose irrelevant passages for training is important for good performance of the model. If we choose too easy negative passages (which have nothing to do with the corresponding relevant passage), the model will quickly learn to distinguish them from positive passages, however, it will work poorly on harder negative passages (for example, those with significant overlap with query or positive passage tokens). Thus, it is common to construct harder negative passages, for example, by selecting irrelevant passages but with high BM25 score. For this tutorial, we will be choosing negative passages randomly. To do it, run corresponding script from [examples/nlp/information_retrieval](https://github.com/NVIDIA/NeMo/blob/main/examples/nlp/information_retrieval):\n",
    "\n",
    "`python construct_random_negatives.py`\n",
    "\n",
    "After this script creates two files query2passages.train.tsv and query2passages.dev.tsv, we will have everything we need to start training."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ZCgWzNBkaQLZ"
   },
   "source": [
    "# Model Training\n",
    "\n",
    "If you have NeMo installed locally, you can also train the model with `examples/nlp/information_retrieval/bert_joint_ir.py`\n",
    "\n",
    "To run training script, use:\n",
    "\n",
    "`python bert_joint_ir.py`"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "GLUE_Benchmark.ipynb",
   "private_outputs": true,
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "metadata": {
     "collapsed": false
    },
    "source": []
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
